package cn.myeasyai.service;

import cn.myeasyai.config.MilvusConfig;
import cn.myeasyai.core.face.ErrorMessage;
import cn.myeasyai.core.face.Face;
import cn.myeasyai.core.face.FaceMessage;
import cn.myeasyai.dto.UserDTO;
import cn.myeasyai.milvus.VectorService;
import cn.myeasyai.util.BusinessException;
import cn.myeasyai.util.IdWorker;
import cn.myeasyai.util.MatrixUtil;
import io.milvus.response.SearchResultsWrapper;
import org.dromara.easyai.entity.ThreeChannelMatrix;
import org.dromara.easyai.matrixTools.Matrix;
import org.dromara.easyai.tools.ImageTools;
import org.dromara.easyai.tools.Picture;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.util.*;

/**
 * @author fushoujiang
 */
@Service
public class FaceService {
    @Resource
    private IdWorker idWorker;
    @Resource
    private Face face;

    @Resource
    private UploadFileService uploadFileService;
    @Resource
    private VectorService vectorService;
    @Resource
    private MilvusConfig milvusConfig;

    public UserDTO faceInit(InputStream inputStream, UserDTO userDTO) {
        ClassPathResource resource = new ClassPathResource("config.xml");

        packageUserDTO(userDTO);
        insertModel(inputStream, userDTO);
        return userDTO;
    }

    public List<UserDTO> faceSearch(InputStream inputStream) {
        final List<Float> floatList = getFloats(inputStream);
        return search(floatList, 5);
    }

    public List<UserDTO> search(List<Float> floatList, Integer topK) {
        final List<SearchResultsWrapper.IDScore> idScoreList = vectorService.search(floatList, topK);
        List<UserDTO> list = new ArrayList<>();
        idScoreList.forEach(idScore -> {
            UserDTO imageDTO = new UserDTO();
            final float score = idScore.getScore();
            final Map<String, Object> fieldValues = idScore.getFieldValues();
            imageDTO.setAutoId(Long.valueOf(String.valueOf( fieldValues.getOrDefault("Auto_id", "-1"))));
            imageDTO.setUserId(Long.valueOf(String.valueOf( fieldValues.getOrDefault("user_id", "-1"))));
            imageDTO.setUserName(String.valueOf((fieldValues.getOrDefault("user_name", ""))));
            imageDTO.setFaceUrl(String.valueOf((fieldValues.getOrDefault("face_url", ""))));
            imageDTO.setFaceFeatureUrl(String.valueOf((fieldValues.getOrDefault("face_feature_url", ""))));
            imageDTO.setScore(Math.sqrt(score));
            list.add(imageDTO);

        });
        return list;
    }

    /**
     * 填充用户信息
     *
     * @param userDTO 用户信息
     */
    private void packageUserDTO(UserDTO userDTO) {
        Long userId = userDTO.getUserId();
        if (Objects.isNull(userId)) {
            userDTO.setUserId(idWorker.nextId());
            final long faceImageId = idWorker.nextId();
            userDTO.setFaceUrlId(faceImageId + "face");
            userDTO.setFaceFeatureUrlId(faceImageId + "face-feature");
            if (Objects.isNull(userDTO.getUserName())){
                userDTO.setUserName("name"+faceImageId);
            }
        }
    }

    private List<Float> getFloats(InputStream inputStream) {
        try {
            ThreeChannelMatrix m = Picture.getThreeMatrix(inputStream, false);
            ErrorMessage errorMessage = face.look(m, idWorker.nextId(), 30);
            if (errorMessage != null) {
                if (errorMessage.getErrorCode() != 0) {
                    throw new BusinessException(errorMessage.getErrorMessage());
                }
                final FaceMessage faceMessage = errorMessage.getFaceMessage();
                final Matrix feature = faceMessage.getFeature();
                return MatrixUtil.matrixToFloatList(feature);
            }
        } catch (Exception e) {
            throw new BusinessException(e.getMessage());
        }
        return Collections.emptyList();
    }


    private void insertModel(final InputStream inputStream, UserDTO userDTO) {
        try {
            final byte[] bytes = inputStream.readAllBytes();
            InputStream copyInputStream = new ByteArrayInputStream(bytes);
            final String faceUrl = uploadFileService.uploadFile(copyInputStream, userDTO.getFaceUrlId());
            InputStream nextBuffer = new ByteArrayInputStream(bytes);

            ThreeChannelMatrix m = Picture.getThreeMatrix(nextBuffer, false);
            ErrorMessage errorMessage = face.look(m, Long.valueOf(userDTO.getUserId()), 30);
            if (errorMessage != null ) {
                if (errorMessage.getErrorCode() != 0){
                    throw new RuntimeException(errorMessage.getErrorMessage());
                }
                final FaceMessage faceMessage = errorMessage.getFaceMessage();
                userDTO.setFaceUrl(faceUrl);
                ThreeChannelMatrix threeChannelMatrix = faceMessage.getChannel();
                InputStream inputStream1 = new ByteArrayInputStream(ImageTools.drawImage(threeChannelMatrix).toByteArray());
                final String url = uploadFileService.uploadFile(inputStream1, userDTO.getFaceFeatureUrlId());
                userDTO.setFaceFeatureUrl(url);
                final Matrix feature = faceMessage.getFeature();
                final List<Float> floats = MatrixUtil.matrixToFloatList(feature);
                vectorService.initUserVector(userDTO, floats);
                return;
            }
            throw new RuntimeException(errorMessage.getErrorMessage());
        } catch (Exception e) {
            e.printStackTrace();
        }
        throw new RuntimeException("初始化图片失败");

    }

}
